from os import listdir, path
import numpy as np
import scipy, cv2, os, sys, argparse, audio
import json, subprocess, random, string
from tqdm import tqdm
import platform
from PIL import Image
import onnxruntime
from queue import Queue
import asyncio


onnxruntime.set_default_logger_severity(3)
from insightface_func.face_detect_crop_single import Face_detect_crop

class LipModel:
    
    def __init__(self, face_path, audio_path, checkpoint_path='checkpoints\wav2lip_256.onnx', outfile='outputs/result_voice.mp4', static=False, 
                 fps=15.0, pads=[0, 10, 0, 0], face_det_batch_size=16, wav2lip_batch_size=1, resize_factor=1, crop=[0, -1, 0, -1], 
                 box=[-1, -1, -1, -1], rotate=False, nosmooth=False):
        self.checkpoint_path = checkpoint_path
        self.face_path = face_path
        self.audio_path = audio_path
        self.outfile = outfile
        self.static = static
        self.fps = fps
        self.pads = pads
        self.face_det_batch_size = face_det_batch_size
        self.wav2lip_batch_size = wav2lip_batch_size
        self.resize_factor = resize_factor
        self.crop = crop
        self.box = box
        self.rotate = rotate
        self.nosmooth = nosmooth
        self.on_complete = False
        
        self.frame_queue = Queue()  # 添加一个队列用于存储视频帧

        if self.checkpoint_path == 'checkpoints\wav2lip_256.onnx' or self.checkpoint_path == 'checkpoints\wav2lip_256_fp16.onnx':
            self.img_size = 256
        else:
            self.img_size = 96
        
        if os.path.isfile(self.face_path) and self.face_path.split('.')[1] in ['jpg', 'png', 'jpeg']:
            self.static = True
        print('self.static'+self.static.__str__())

    def load_model(self, device):
        model_path = self.checkpoint_path
        session_options = onnxruntime.SessionOptions()
        session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
        providers = ["CPUExecutionProvider"]
        if device == 'cuda':
            providers = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}), "CPUExecutionProvider"]
        session = onnxruntime.InferenceSession(model_path, sess_options=session_options, providers=providers)
        return session

    def get_smoothened_boxes(self, boxes, T):
        for i in range(len(boxes)):
            if i + T > len(boxes):
                window = boxes[len(boxes) - T:]
            else:
                window = boxes[i: i + T]
            boxes[i] = np.mean(window, axis=0)
        return boxes

    def face_detect(self, images):
        detector = Face_detect_crop(name='antelope', root='./insightface_func/models')
        detector.prepare(ctx_id=0, det_thresh=0.3, det_size=(320, 320), mode='none')

        predictions = []
        for i in tqdm(range(0, len(images))):
            bbox = detector.getBox(images[i])
            predictions.append(bbox)

        results = []
        pady1, pady2, padx1, padx2 = self.pads
        for rect, image in zip(predictions, images):
            if rect is None:
                cv2.imwrite('temp/faulty_frame.jpg', image)
                raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')

            y1 = max(0, rect[1] - pady1)
            y2 = min(image.shape[0], rect[3] + pady2)
            x1 = max(0, rect[0] - padx1)
            x2 = min(image.shape[1], rect[2] + padx2)

            results.append([x1, y1, x2, y2])

        boxes = np.array(results)
        if not self.nosmooth:
            boxes = self.get_smoothened_boxes(boxes, T=5)
        results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]

        del detector
        return results

    def datagen(self, frames, mels):
        img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []

        if self.box[0] == -1:
            if not self.static:
                face_det_results = self.face_detect(frames)
            else:
                face_det_results = self.face_detect([frames[0]])
        else:
            print('Using the specified bounding box instead of face detection...')
            y1, y2, x1, x2 = self.box
            face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]

        for i, m in enumerate(mels):
            idx = 0 if self.static else i % len(frames)
            frame_to_save = frames[idx].copy()
            face, coords = face_det_results[idx].copy()

            face = cv2.resize(face, (self.img_size, self.img_size))

            img_batch.append(face)
            mel_batch.append(m)
            frame_batch.append(frame_to_save)
            coords_batch.append(coords)

            if len(img_batch) >= self.wav2lip_batch_size:
                img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

                img_masked = img_batch.copy()
                img_masked[:, self.img_size // 2:] = 0

                img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
                mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])

                yield img_batch, mel_batch, frame_batch, coords_batch
                img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []

        if len(img_batch) > 0:
            img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

            img_masked = img_batch.copy()
            img_masked[:, self.img_size // 2:] = 0

            img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
            mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])

            yield img_batch, mel_batch, frame_batch, coords_batch

    def process(self):
        im = cv2.imread(self.face_path)
        if not os.path.isfile(self.face_path):
            raise ValueError('--face argument must be a valid path to video/image file')

        elif self.face_path.split('.')[1] in ['jpg', 'png', 'jpeg']:
            full_frames = [cv2.imread(self.face_path)]
            fps = self.fps

        else:
            video_stream = cv2.VideoCapture(self.face_path)
            fps = video_stream.get(cv2.CAP_PROP_FPS)

            print('Reading video frames...')

            full_frames = []
            while 1:
                still_reading, frame = video_stream.read()
                if not still_reading:
                    video_stream.release()
                    break
                if self.resize_factor > 1:
                    frame = cv2.resize(frame, (frame.shape[1] // self.resize_factor, frame.shape[0] // self.resize_factor))

                if self.rotate:
                    frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)

                y1, y2, x1, x2 = self.crop
                if x2 == -1:
                    x2 = frame.shape[1]
                if y2 == -1:
                    y2 = frame.shape[0]

                frame = frame[y1:y2, x1:x2]

                full_frames.append(frame)

        print("Number of frames available for inference: " + str(len(full_frames)))

        if not self.audio_path.endswith('.wav'):
            print('Extracting raw audio...')
            command = 'ffmpeg -y -i {} -strict -2 {}'.format(self.audio_path, 'temp/temp.wav')

            subprocess.call(command, shell=True)
            self.audio_path = 'temp/temp.wav'

        wav = audio.load_wav(self.audio_path, 16000)
        mel = audio.melspectrogram(wav)

        if np.isnan(mel.reshape(-1)).sum() > 0:
            raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')

        mel_chunks = []
        mel_idx_multiplier = 80. / fps
        i = 0
        while 1:
            start_idx = int(i * mel_idx_multiplier)
            if start_idx + 16 > len(mel[0]):
                mel_chunks.append(mel[:, len(mel[0]) - 16:])
                break
            mel_chunks.append(mel[:, start_idx: start_idx + 16])
            i += 1

        print("Length of mel chunks: {}".format(len(mel_chunks)))

        full_frames = full_frames[:len(mel_chunks)]

        batch_size = 1
        gen = self.datagen(full_frames.copy(), mel_chunks)

        device = 'cuda' if onnxruntime.get_device() == 'GPU' else 'cpu'

        for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, total=int(np.ceil(float(len(mel_chunks)) / batch_size)))):
            if i == 0:
                model = self.load_model(device)
                frame_h, frame_w = full_frames[0].shape[:-1]
                out = cv2.VideoWriter('temp/result.avi', cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))

            img_batch = img_batch.transpose((0, 3, 1, 2)).astype(np.float32)
            mel_batch = mel_batch.transpose((0, 3, 1, 2)).astype(np.float32)

            pred = model.run(None, {'mel_spectrogram': mel_batch, 'video_frames': img_batch})[0][0]
            pred = pred.transpose(1, 2, 0) * 255
            pred = pred.astype(np.uint8)
            pred = pred.reshape((1, self.img_size, self.img_size, 3))

            for p, f, c in zip(pred, frames, coords):
                y1, y2, x1, x2 = c
                p = cv2.resize(p, (x2 - x1, y2 - y1))
                f[y1:y2, x1:x2] = p
                out.write(f)
                cv2.imshow("Result", f)
                cv2.waitKey(1)
                
                # 将视频帧放入队列
                self.frame_queue.put(f)  
                # print(f"Processing finished for {self.face_path} and {self.audio_path}.")
                

        out.release()
        # 调用回调函数
        self.on_complete = True

        command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(self.audio_path, 'temp/result.avi', self.outfile)
        subprocess.call(command, shell=platform.system() != 'Windows')
        os.remove('temp/result.avi')

# 使用示例
if __name__ == '__main__':
    processor = LipModel(face_path='D:\some.mp4',audio_path='D:\some.wav')
    processor.process()